
import torch.nn as nn 

from .layers import TwoConv, Upsample, Conv
import warnings

import torch.nn as nn
import torch.nn.functional as F
# Copyright (c) OpenMMLab. All rights reserved.
import torch
import torch.nn as nn


class FPN(nn.Module):
    def __init__(self, features, out_channels, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.features = features
        
        self.convs = []
        for f in self.features:
            self.convs.append(TwoConv(f, out_channels))
        self.convs = nn.ModuleList(self.convs)

        # self.ups = []
        # for i in range(stage):
        #     self.ups.append(Upsample(2*out_channels, out_channels))

    def forward(self, multi_scale_features):

        # for f in multi_scale_features:
        # f_bottom = multi_scale_features[0]

        outputs = []
        index = 0
        for f in multi_scale_features:
            outputs.append(self.convs[index](f))
            index += 1
        return outputs
    


def resize(input,
           size=None,
           scale_factor=None,
           mode='nearest',
           align_corners=None,
           warning=True):
    if warning:
        if size is not None and align_corners:
            input_h, input_w = tuple(int(x) for x in input.shape[2:])
            output_h, output_w = tuple(int(x) for x in size)
            if output_h > input_h or output_w > input_w:
                if ((output_h > 1 and output_w > 1 and input_h > 1
                     and input_w > 1) and (output_h - 1) % (input_h - 1)
                        and (output_w - 1) % (input_w - 1)):
                    warnings.warn(
                        f'When align_corners={align_corners}, '
                        'the output would more aligned if '
                        f'input size {(input_h, input_w)} is `x+1` and '
                        f'out size {(output_h, output_w)} is `nx+1`')
    return F.interpolate(input, size, scale_factor, mode, align_corners)

class MultiStageMerging(nn.Module):

    def __init__(self,
                 in_channels,
                 out_channels,
                 kernel_size=1,
                 align_corners=False,
                 ):
        super(MultiStageMerging, self).__init__()
        assert isinstance(in_channels, list)
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.align_corners = align_corners
        # self.down = ConvModule(
        #     sum(in_channels),
        #     out_channels,
        #     kernel_size,
        #     padding=0,
        #     conv_cfg=conv_cfg,
        #     norm_cfg=norm_cfg,
        #     act_cfg=act_cfg
        # )
        self.down = Conv(sum(in_channels), out_channels, kernel_size, stride=1, padding=0)
        

    def forward(self, inputs):
        assert len(inputs) == len(self.in_channels)
        outs = list()
        size = inputs[0].shape[2:]
        for index, input in enumerate(inputs):
            input = resize(input,
                           size=size,
                           mode='bilinear',
                           align_corners=self.align_corners)
            outs.append(input)
        out = torch.cat(outs, dim=1)
        out = self.down(out)
        return out